package org.emailclassifier.preprocessing;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.Random;
import java.util.Set;

public class EmailDataSet {

	
	private HashMap<String, ArrayList<DataSetDocument>> m_documents = new HashMap<String, ArrayList<DataSetDocument>>();

	private HashMap<String, Integer> m_globalTermFrequency = new HashMap<String, Integer>();
	private HashMap<String, Integer> m_documentFrequency = new HashMap<String, Integer>();
	private HashMap<String, HashMap<String, Integer>> m_categoryTermFrequency = new HashMap<String, HashMap<String, Integer>>();
	
	private boolean m_stoplist = false; 
	
	public HashMap<String, Integer> getGlobalTermFrequency()
	{
		return m_globalTermFrequency;
	}
	
	public void addDocument( String category, String body)
	{
		DataSetDocument doc = new DataSetDocument(body, m_stoplist);
		if (m_documents.containsKey(category)) {
			ArrayList<DataSetDocument> list = m_documents.get(category);
			list.add(doc);
		} else {
			ArrayList<DataSetDocument> list = new ArrayList<DataSetDocument>();
			list.add(doc);
			m_documents.put(category, list);
		}
		
		updateFrequencies(doc, category);
	}
	
	public void addDocument( String category, DataSetDocument doc )
	{
		if (m_documents.containsKey(category)) {
			ArrayList<DataSetDocument> list = m_documents.get(category);
			list.add(doc);
		} else {
			ArrayList<DataSetDocument> list = new ArrayList<DataSetDocument>();
			list.add(doc);
			m_documents.put(category, list);
		}
		
		updateFrequencies(doc, category);	
	}
	
	private void updateFrequencies(DataSetDocument doc, String category)
	{
		HashMap<String, Integer> freq = doc.getTermFrequencies();
		
		HashMap<String, Integer> catfreq;
		if (m_categoryTermFrequency.containsKey(category)) {
			catfreq = m_categoryTermFrequency.get(category);
		} else {
			catfreq = new HashMap<String, Integer>();
		}
		
				
		for (String term : freq.keySet())
		{
			if (m_documentFrequency.containsKey(term)) {
				m_documentFrequency.put(term, m_documentFrequency.get(term) + 1);
			} else {
				m_documentFrequency.put(term, 1);
			}
			
			if (m_globalTermFrequency.containsKey(term)) {
				m_globalTermFrequency.put(term, m_globalTermFrequency.get(term) + freq.get(term));
			} else {
				m_globalTermFrequency.put(term, freq.get(term));
			}
			
			if (catfreq.containsKey(term)) {
				catfreq.put(term, catfreq.get(term) + freq.get(term));
			} else {
				catfreq.put(term, freq.get(term));
			}
		}
		
		m_categoryTermFrequency.put(category, catfreq);
	}
	
	public Set<String> getGlobalVocabulary()
	{
		return m_documentFrequency.keySet();
	}
	
	public Set<String> getCategoryVocabulary(String category)
	{
		if (m_categoryTermFrequency.containsKey(category))
		{
			return m_categoryTermFrequency.get(category).keySet();
		}
		return null;
	}

	public int getGlobalTermFrequency(String term)
	{
		if (m_globalTermFrequency.containsKey(term)) {
			return m_globalTermFrequency.get(term);
		}
		return 0;
	}
	
	public int getCategoryTermFrequency(String category, String term)
	{
		if (m_categoryTermFrequency.containsKey(category)) {
			HashMap<String, Integer> cat = m_categoryTermFrequency.get(category);
			if (cat.containsKey(term)) {
				return cat.get(term);
			}
			return 0;
				
		}
		return 0;
	}
	
	public HashMap<String, Integer> getCategoryTermFrequencies(String category)
	{
		if (m_categoryTermFrequency.containsKey(category)) {
			return m_categoryTermFrequency.get(category);
		}
		return new HashMap<String, Integer>();
	}
	
	public int getDocumentFrequency(String term)
	{
		if (m_documentFrequency.containsKey(term)) {
			return m_documentFrequency.get(term);
		}
		return 0;
	}
	
	public int getGlobalDocumentCount()
	{
		int qty = 0;
		for (String category : m_documents.keySet())
		{
			qty += m_documents.get(category).size();
		}
		return qty;
	}
	
	public int getCategoryCount()
	{
		return m_documents.keySet().size();
	}
	
	public Set<String> getCategories()
	{
		return m_documents.keySet();		
	}
	
	public int getDocumentsInCategoryCount(String category)
	{
		if (m_documents.containsKey(category))
		{
			return m_documents.get(category).size();
		}
		return 0;
	}
	
	public DataSetDocument[] getDocumentsInCategory(String category)
	{
		if (m_documents.containsKey(category))
		{
			return m_documents.get(category).toArray(new DataSetDocument[0]);
		}
		return new DataSetDocument[0];
	}
	
	public EmailDataSet[] makeCrossValidationSet(double testFrac)
	{
		EmailDataSet train = new EmailDataSet();
		EmailDataSet test = new EmailDataSet();
		
		for (String category : m_documents.keySet())
		{
			ArrayList<DataSetDocument> docs = m_documents.get(category);
			
			// Randomize
			Collections.shuffle(docs, new Random(System.nanoTime()));
			
			int trainCutoff = (int)((double)docs.size() * testFrac);
			int i;
			
			if (trainCutoff < 1 && docs.size() > 0 )
				trainCutoff = 1;
			
			for (i = 0; i < trainCutoff; i++)
			{
				train.addDocument(category, docs.get(i));
			}
			
			for (; i < docs.size(); i++)
			{
				test.addDocument(category, docs.get(i));
			}
		}
		
		EmailDataSet[] array = new EmailDataSet[2];
		array[0] = train;
		array[1] = test;
		
		return array;
	}

	public boolean getStoplist() {
		return m_stoplist;
	}

	public void setStoplist(boolean stoplist) {
		m_stoplist = stoplist;
	}
}
